# 主要看index能不能配合模型精确地召回自己原有的数据，且距离是0
import sys
import torch
sys.path.append('..')
from sent2vec import Sent2VecEmbeddings
import json
import numpy as np
# BAAI/bge-large-zh-v1.5


with open('/home/lxy/wxbdata/db/merge.json','r',encoding='utf-8') as f:
    data=json.load(f)
corpus=[d['question'] for d in data]
print(len(corpus))
maxNorm=torch.load('/data/lxy/RAT/gte_11m21d_test_accelerate/maxNorm.pt')

import faiss
import time

print('要开始了')
model=Sent2VecEmbeddings(model_name='BAAI/bge-large-zh-v1.5')
print('模型加载完了')
start_time = time.time()
acc=0
index=faiss.read_index('/data/lxy/RAT/gte_11m21d_test_accelerate/index.faiss')
for q in corpus:
    query_embedding = model.embed_query(q).to('cpu').unsqueeze(0)/maxNorm

    # vector = np.array(embeddings, dtype=np.float32)
    #         if normalize_L2:
    #             faiss.normalize_L2(vector)
    # query_embedding=np.array(query_embedding, dtype=np.float32)
    # faiss.normalize_L2(query_embedding)
    # 邱吉雨有没有指导过学生？
    D,I=index.search(query_embedding,k=5)
    # print(query_embedding)
    # print(D,I)
    if acc%500==0:
        print(acc)
    if corpus[int(I[0][0])]==q and D[0][0]<0.05:
        acc+=1

    else:
        for i in I[0]:
            print(f'D: {D[0][0]} {corpus[int(i)]}')

end_time = time.time()
print("耗时: {:.2f}秒".format(end_time - start_time))
print(f'完备率={acc/len(corpus)}')


# 这上面要是成功了，说明索引本身是没问题的。 成功了？

# corpus_embeddings = torch.load('/home/ruanjh/DPR/qq_bge_2wfinetuned_epoch10_4w_idxflatl2.pt')

# # Query sentences:
# queries = [
#     '你是谁研发的大模型？',
#     '我想上东大！',
#     '怎样才能知道自己的档案是否已经被审核通过？',
#     '是否能够在提前批阶段同时选择军校和警校？']

# for query in queries:
#     query_embedding = model.embed_query(query)
#     hits = semantic_search(query_embedding, corpus_embeddings, top_k=5)
#     print("\n\n======================\n\n")
#     print("Query:", query)
#     print("\nTop 5 most similar sentences in corpus:")
#     hits = hits[0]  # Get the hits for the first query
#     for hit in hits:
#         print(corpus[hit['corpus_id']], "(Score: {:.4f})".format(hit['score']))